/* Wang Wen(Albert Wang) jawangwen@gmail.com */

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <memory.h>
#ifdef _WIN32
#include <windows.h>
#else
#include <sys/mman.h>
#include <unistd.h>
#include <errno.h>
#endif
#include "rt_stub.h"

#ifdef __x86_64__
#define CODE_MAX        8
#else
#define CODE_MAX        5
#endif

#define AVAILABLE       1
#define NOT_AVAILABLE   0
#define ENABLE          1
#define NOT_ENABLE      0
#define WRITABLE        1
#define NOT_WRITABLE    0

#define INFO_ENTRY(info_head_ptr) \
        ((stub_info *)((char *)(info_head_ptr)-(unsigned long)(&((stub_info *)0)->head)))

#define INFO_HEAD_EMPTY(info_head_ptr)   ((info_head_ptr) == &g_info_sets.head)
#define INFO_HEAD_FIRST()                (g_info_sets.head.next)

#define FIRST_INFO_HEAD(info_head_ptr)   ((info_head_ptr) = g_info_sets.head.next)
#define NEXT_INFO_HEAD(info_head_ptr)    ((info_head_ptr) = (info_head_ptr)->next)

#define INIT_INFO_HEAD(info_head_ptr) do {   \
    (info_head_ptr)->next = (info_head_ptr); \
    (info_head_ptr)->prev = (info_head_ptr); \
} while (0)

#define INSERT_INFO_HEAD(info_head_ptr) do {             \
    (info_head_ptr)->next     = INFO_HEAD_FIRST();       \
    (info_head_ptr)->prev     = INFO_HEAD_FIRST()->prev; \
    INFO_HEAD_FIRST()->prev   = (info_head_ptr);         \
    INFO_HEAD_FIRST()         = (info_head_ptr);         \
} while (0)

#define DELETE_INFO_HEAD(info_head_ptr) do {             \
    (info_head_ptr)->next->prev = (info_head_ptr)->prev; \
    (info_head_ptr)->prev->next = (info_head_ptr)->next; \
    INIT_INFO_HEAD(info_head_ptr);                       \
} while (0)

#define FOREACH_INFO_HEAD(info_head_ptr) for ( \
    FIRST_INFO_HEAD(info_head_ptr);            \
    !INFO_HEAD_EMPTY(info_head_ptr);           \
    NEXT_INFO_HEAD(info_head_ptr)              \
)

#define FOREACH_CODE(byte_ptr, func_addr) for (                 \
    (byte_ptr) = (unsigned char *)(func_addr);                  \
    (byte_ptr) - (unsigned char *)(func_addr) < (CODE_MAX - 1); \
    (byte_ptr)++                                                \
)

#define CODE_IS_RET(byte)           ((unsigned char)(byte) == 0xC3)
#define CODE_INDEX(new_code)        (((unsigned char *)new_code) + g_std_code_index)

#ifdef __x86_64__
#define GET_ADDR(info_ptr)          ((unsigned int)((info_ptr)->stub))
#else
#define GET_ADDR(info_ptr)          GET_RELATIVE_ADDR((info_ptr)->stub, (info_ptr)->func)
#define GET_RELATIVE_ADDR(dst, src) ((dst)- (src) - CODE_MAX)
#endif

typedef struct stub_info_head {
    struct stub_info_head * prev;
    struct stub_info_head * next;
} stub_info_head;

#pragma pack(1)
typedef struct stub_info {
    stub_info_head   head;
    unsigned int     enable;
    void           * func;
    void           * stub;
	unsigned char    old_code[CODE_MAX];
#ifdef _WIN32
    DWORD            old_protect;
#endif
} stub_info;
#pragma pack()

/*****************************************************************************
* SYS32 : __asm("jmp  $mock")                  : 0xE9 {$mock -$orig-5}       *
* SYS32 : __asm("movl $mock, %eax; jmpl %eax") : 0xB8 {$mock} 0xFF 0xE0      *
* SYS64 : __asm("movq $mock, %rax; jmpq %rax") : 0x48 0xB8 {$mock} 0xFF 0xE0 * 
*****************************************************************************/
#ifdef __x86_64__
static const unsigned char g_std_code[CODE_MAX] = {0x48,0xB8, 0x0,0x0,0x0,0x0, 0xFF,0xE0};
static const unsigned int  g_std_code_index = 2;
#else
static const unsigned char g_std_code[CODE_MAX] = {0xE9, 0x0,0x0,0x0,0x0};
static const unsigned int  g_std_code_index = 1;
#endif

static stub_info g_info_sets = {{&g_info_sets.head, &g_info_sets.head}};

static __inline__ void insert_info (stub_info * info_ptr);
static __inline__ void delete_info (stub_info * info_ptr);

static __inline__ void save_code (stub_info * info_ptr)
{
    memcpy((void *)info_ptr->old_code, (void *)info_ptr->func, CODE_MAX);
}

static __inline__ void make_code (stub_info     * info_ptr,
                                  unsigned char   new_code[CODE_MAX])
{
    unsigned char * code_index;
    int             addr;

    code_index = CODE_INDEX(new_code);
    addr       = GET_ADDR(info_ptr);
    memcpy(new_code,   g_std_code, CODE_MAX);
    memcpy(code_index, &addr,    sizeof(void*));
}

static __inline__ void enable_code (stub_info * info_ptr)
{
    unsigned char new_code[CODE_MAX];

    make_code(info_ptr, new_code);
    memcpy((void *)info_ptr->func, (void *)new_code, CODE_MAX);
    insert_info(info_ptr);
}

static __inline__ void disable_code (stub_info * info_ptr)
{
    if (info_ptr->enable) {
        memcpy((void *)info_ptr->func, (void *)info_ptr->old_code, CODE_MAX);
        delete_info(info_ptr);
    }
}

static __inline__ stub_info * find_info (void * func)
{
    stub_info_head * info_head_ptr;
    stub_info      * info_ptr;

    FOREACH_INFO_HEAD(info_head_ptr) {
        info_ptr = INFO_ENTRY(info_head_ptr);
        if (info_ptr->func == func) {
            return info_ptr;
        }
    }
    return NULL;
}

static __inline__ void insert_info (stub_info * info_ptr)
{
    if (!info_ptr->enable) {
        INSERT_INFO_HEAD(&info_ptr->head);
        info_ptr->enable = ENABLE;
    }
}

static __inline__ void delete_info (stub_info * info_ptr)
{
    if (info_ptr->enable) {
        DELETE_INFO_HEAD(&info_ptr->head);
        info_ptr->enable = NOT_ENABLE;
    }
}

static __inline__ stub_info * alloc_info (void * func, void * stub)
{
    stub_info * info_ptr;

    info_ptr = (stub_info *)malloc(sizeof(stub_info));
    if (info_ptr == NULL) {
        return NULL;
    }

    memset((void *)info_ptr, 0, sizeof(stub_info));
    INIT_INFO_HEAD(&info_ptr->head);
    info_ptr->enable = NOT_ENABLE;
    info_ptr->func   = func;
    info_ptr->stub   = stub;
    save_code(info_ptr);

    return info_ptr;
}

static __inline__ void free_info (stub_info * info_ptr)
{
    if (info_ptr->enable) {
        delete_info(info_ptr);
    }
    free(info_ptr);
}

static __inline__ int alloc_code_writable (stub_info * info_ptr)
{
    int writable = WRITABLE;
#ifdef _WIN32
    RT_STUB_TRACE("VirtualProtect: alloc for writing code");
    if (!VirtualProtect(info_ptr->func, CODE_MAX, PAGE_WRITECOPY, &info_ptr->old_protect)) {
        writable = NOT_WRITABLE;
    }
#else
    int pagesize = sysconf(_SC_PAGE_SIZE);
    /* XXX: how can we save the old prot */
    RT_STUB_TRACE("mprotect: alloc for writing code");
    if(mprotect((void*)((unsigned long)info_ptr->func & (~(pagesize - 1))),
                pagesize,
                PROT_READ | PROT_WRITE | PROT_EXEC)) {
        RT_STUB_ERROR("but fail: %s", strerror(errno));
        writable = NOT_WRITABLE;
    }
#endif
    return writable;
}

static __inline__ void free_code_writable (stub_info * info_ptr)
{
#ifdef _WIN32
    RT_STUB_TRACE("VirtualProtect: free for writing code");
    if (!VirtualProtect(info_ptr->func, 
                        CODE_MAX,
                        info_ptr->old_protect, 
                        &info_ptr->old_protect)) {
        RT_STUB_WARN("but fail!");
    }
#endif
}

static __inline__ int code_available (void * func)
{
    unsigned char * bytePtr;

    FOREACH_CODE(bytePtr, func) {
        if (CODE_IS_RET(*bytePtr)) {
            return NOT_AVAILABLE;
        }
    }
    return AVAILABLE;
}

int set_stub (void * func, void * stub)
{
    stub_info * info_ptr;
    int         writable;
    void      * old_stub;

    if (func == NULL) {
        RT_STUB_ERROR("func<NULL> invalid!");
        return -1;
    }
    if (stub == NULL) {
        RT_STUB_ERROR("stub<NULL> invalid!");
        return -1;
    }

    info_ptr = find_info(func);
    if (info_ptr == NULL) {
        if (!code_available(func)) {
            RT_STUB_ERROR("func<0x%x> cannot be stubed!", (unsigned int)func);
            return -1;
        }
        info_ptr = alloc_info(func, stub);
        if (info_ptr == NULL) {
            RT_STUB_ERROR("not enough code to save stub_info for func<0x%x>!",
                         (unsigned int)func);
            return -1;
        }
    }

    old_stub = info_ptr->stub;
    if (stub != old_stub) {
        RT_STUB_WARN("stub covered from func<0x%x> to func<0x%x>!",
                     (unsigned int)old_stub,
                     (unsigned int)stub);
        info_ptr->stub = stub;
    }

    writable = alloc_code_writable(info_ptr);
    if (!writable) {
        RT_STUB_ERROR("fail to alloc for writing code!");
        info_ptr->stub = old_stub;
        if (!info_ptr->enable) {
            free_info(info_ptr);
        }
        return -1;
    }
    enable_code(info_ptr);
    free_code_writable(info_ptr);

    return 0;
}

int del_stub (void * func)
{
    stub_info * info_ptr;
    int         writable;

    if (func == NULL) {
        RT_STUB_ERROR("func<NULL> invalid!");
        return -1;
    }

    info_ptr = find_info(func);
    if (info_ptr == NULL) {
        return 0;
    }

    writable = alloc_code_writable(info_ptr);
    if(!writable) {
        RT_STUB_ERROR("fail to alloc for writing code!");
        return -1;
    }
    disable_code(info_ptr);
    free_code_writable(info_ptr);

    return 0;
}

int clean_stub (void)
{
    stub_info_head * info_head_ptr;
    stub_info      * info_ptr;
    int              success;

    success = 0;
    FIRST_INFO_HEAD(info_head_ptr);
    while (!INFO_HEAD_EMPTY(info_head_ptr)) {
        info_ptr = INFO_ENTRY(info_head_ptr);
        RT_STUB_TRACE("[del_stub] <0x%x>", (unsigned int)info_ptr->func);
        if (del_stub(info_ptr->func)) {
            if (find_info(info_ptr->func) == info_ptr) {
                DELETE_INFO_HEAD(info_head_ptr);
                free(info_ptr);
            }
            success = -1;
        }
        FIRST_INFO_HEAD(info_head_ptr);
    }

    return success;
}

